// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/common/base/exceptions.h>

#include <melon/random.h>

#include <cstdint>
#include <optional>
#include <random>

namespace kumo::pollux::random {
    // Set a custom seed to be returned in all following getSeed() calls.
    // Should be only used in unit tests.
    void setSeed(uint32_t);

    // Return a true random seed unless setSeed() is called before.
    uint32_t getSeed();

    /// Utility class to accelerate random sampling based on Bernoulli trials.
    /// Internally this keeps the number of skips for next hit.  User can consume a
    /// bulk of trials calling the `nextSkip' then `consume', or call `testOne` to
    /// do the trials one by one.
    class RandomSkipTracker {
    public:
        explicit RandomSkipTracker(double sampleRate);

        RandomSkipTracker(const RandomSkipTracker &) = delete;

        RandomSkipTracker &operator=(const RandomSkipTracker &) = delete;

        /// Return the number of skips need to get a hit.  Must be called before
        /// calling `consume'.
        uint64_t nextSkip() {
            if (sampleRate_ == 0) {
                return std::numeric_limits<uint64_t>::max();
            }
            if (skip_.has_value()) {
                return *skip_;
            }
            skip_ = dist_(rng_);
            return *skip_;
        }

        /// Consume the remaining skips followed by at most one hit.
        void consume(uint64_t numElements) {
            if (sampleRate_ == 0) {
                return;
            }
            POLLUX_DCHECK(skip_.has_value());
            if (*skip_ >= numElements) {
                *skip_ -= numElements;
            } else {
                POLLUX_DCHECK_EQ(numElements - *skip_, 1);
                skip_.reset();
            }
        }

        /// Consume one trial and return the result.
        bool testOne() {
            if (sampleRate_ == 0) {
                return false;
            }
            if (nextSkip() == 0) {
                skip_.reset();
                return true;
            }
            --*skip_;
            return false;
        }

        double sampleRate() const {
            return sampleRate_;
        }

    private:
        const double sampleRate_;
        std::geometric_distribution<uint64_t> dist_;
        melon::Random::DefaultGenerator rng_;
        std::optional<uint64_t> skip_;
    };
} // namespace kumo::pollux::random
